{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Saving PyBaMM models to file"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Models which are discretised (i.e. ready to solve/ previously solved, see [this notebook](https://github.com/pybamm-team/PyBaMM/blob/develop/docs/source/examples/notebooks/spatial_methods/finite-volumes.ipynb) for more information on the pybamm.Discretisation class) can be serialised and saved to a JSON file, ready to be read in again either in PyBaMM, or a different modelling library. \n",
    "\n",
    "In the example below, we build a basic DFN model, and then save the model out to `sim_model_example.json`, which should have appear in the 'models' directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install pybamm -q    # install PyBaMM if it is not installed\n",
    "import pybamm\n",
    "\n",
    "# do the example\n",
    "dfn_model = pybamm.lithium_ion.DFN()\n",
    "dfn_sim = pybamm.Simulation(dfn_model)\n",
    "# discretise and build the model\n",
    "dfn_sim.build()\n",
    "\n",
    "dfn_sim.save_model(\"sim_model_example\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This model file can then be read in and solved by choosing a solver, and running as below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<pybamm.solvers.solution.Solution at 0x2a4e10550>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Recreate the pybamm model from the JSON file\n",
    "new_dfn_model = pybamm.load_model(\"sim_model_example.json\")\n",
    "\n",
    "sim_reloaded = pybamm.Simulation(new_dfn_model)\n",
    "sim_reloaded.solve([0, 3600])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It would be nice to plot the results of the two models, to confirm that they are producing the same result.\n",
    "\n",
    "However, notice that running the code below generates an error stating that the model variables were not provided during the reading in of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "No variables to plot",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mAttributeError\u001B[0m                            Traceback (most recent call last)",
      "\u001B[1;32m/Users/pipliggins/Documents/repos/pybamm-local/docs/source/examples/notebooks/models/saving_models.ipynb Cell 7\u001B[0m line \u001B[0;36m8\n\u001B[1;32m      <a href='vscode-notebook-cell:/Users/pipliggins/Documents/repos/pybamm-local/docs/source/examples/notebooks/models/saving_models.ipynb#W6sZmlsZQ%3D%3D?line=4'>5</a>\u001B[0m     plot_sim\u001B[39m.\u001B[39msolve([\u001B[39m0\u001B[39m, \u001B[39m3600\u001B[39m])\n\u001B[1;32m      <a href='vscode-notebook-cell:/Users/pipliggins/Documents/repos/pybamm-local/docs/source/examples/notebooks/models/saving_models.ipynb#W6sZmlsZQ%3D%3D?line=5'>6</a>\u001B[0m     sims\u001B[39m.\u001B[39mappend(plot_sim)\n\u001B[0;32m----> <a href='vscode-notebook-cell:/Users/pipliggins/Documents/repos/pybamm-local/docs/source/examples/notebooks/models/saving_models.ipynb#W6sZmlsZQ%3D%3D?line=7'>8</a>\u001B[0m pybamm\u001B[39m.\u001B[39;49mdynamic_plot(sims, time_unit\u001B[39m=\u001B[39;49m\u001B[39m\"\u001B[39;49m\u001B[39mseconds\u001B[39;49m\u001B[39m\"\u001B[39;49m)\n",
      "File \u001B[0;32m~/Documents/repos/pybamm-local/pybamm/plotting/dynamic_plot.py:20\u001B[0m, in \u001B[0;36mdynamic_plot\u001B[0;34m(*args, **kwargs)\u001B[0m\n\u001B[1;32m      8\u001B[0m \u001B[39m\u001B[39m\u001B[39m\"\"\"\u001B[39;00m\n\u001B[1;32m      9\u001B[0m \u001B[39mCreates a :class:`pybamm.QuickPlot` object (with arguments 'args' and keyword\u001B[39;00m\n\u001B[1;32m     10\u001B[0m \u001B[39marguments 'kwargs') and then calls :meth:`pybamm.QuickPlot.dynamic_plot`.\u001B[39;00m\n\u001B[0;32m   (...)\u001B[0m\n\u001B[1;32m     17\u001B[0m \u001B[39m    The 'QuickPlot' object that was created\u001B[39;00m\n\u001B[1;32m     18\u001B[0m \u001B[39m\"\"\"\u001B[39;00m\n\u001B[1;32m     19\u001B[0m kwargs_for_class \u001B[39m=\u001B[39m {k: v \u001B[39mfor\u001B[39;00m k, v \u001B[39min\u001B[39;00m kwargs\u001B[39m.\u001B[39mitems() \u001B[39mif\u001B[39;00m k \u001B[39m!=\u001B[39m \u001B[39m\"\u001B[39m\u001B[39mtesting\u001B[39m\u001B[39m\"\u001B[39m}\n\u001B[0;32m---> 20\u001B[0m plot \u001B[39m=\u001B[39m pybamm\u001B[39m.\u001B[39;49mQuickPlot(\u001B[39m*\u001B[39;49margs, \u001B[39m*\u001B[39;49m\u001B[39m*\u001B[39;49mkwargs_for_class)\n\u001B[1;32m     21\u001B[0m plot\u001B[39m.\u001B[39mdynamic_plot(kwargs\u001B[39m.\u001B[39mget(\u001B[39m\"\u001B[39m\u001B[39mtesting\u001B[39m\u001B[39m\"\u001B[39m, \u001B[39mFalse\u001B[39;00m))\n\u001B[1;32m     22\u001B[0m \u001B[39mreturn\u001B[39;00m plot\n",
      "File \u001B[0;32m~/Documents/repos/pybamm-local/pybamm/plotting/quick_plot.py:146\u001B[0m, in \u001B[0;36mQuickPlot.__init__\u001B[0;34m(self, solutions, output_variables, labels, colors, linestyles, shading, figsize, n_rows, time_unit, spatial_unit, variable_limits)\u001B[0m\n\u001B[1;32m    144\u001B[0m \u001B[39m# check variables have been provided after any serialisation\u001B[39;00m\n\u001B[1;32m    145\u001B[0m \u001B[39mif\u001B[39;00m \u001B[39many\u001B[39m(\u001B[39mlen\u001B[39m(m\u001B[39m.\u001B[39mvariables) \u001B[39m==\u001B[39m \u001B[39m0\u001B[39m \u001B[39mfor\u001B[39;00m m \u001B[39min\u001B[39;00m models):\n\u001B[0;32m--> 146\u001B[0m     \u001B[39mraise\u001B[39;00m \u001B[39mAttributeError\u001B[39;00m(\u001B[39m\"\u001B[39m\u001B[39mNo variables to plot\u001B[39m\u001B[39m\"\u001B[39m)\n\u001B[1;32m    148\u001B[0m \u001B[39mself\u001B[39m\u001B[39m.\u001B[39mn_rows \u001B[39m=\u001B[39m n_rows \u001B[39mor\u001B[39;00m \u001B[39mint\u001B[39m(\n\u001B[1;32m    149\u001B[0m     \u001B[39mlen\u001B[39m(output_variables) \u001B[39m/\u001B[39m\u001B[39m/\u001B[39m np\u001B[39m.\u001B[39msqrt(\u001B[39mlen\u001B[39m(output_variables))\n\u001B[1;32m    150\u001B[0m )\n\u001B[1;32m    151\u001B[0m \u001B[39mself\u001B[39m\u001B[39m.\u001B[39mn_cols \u001B[39m=\u001B[39m \u001B[39mint\u001B[39m(np\u001B[39m.\u001B[39mceil(\u001B[39mlen\u001B[39m(output_variables) \u001B[39m/\u001B[39m \u001B[39mself\u001B[39m\u001B[39m.\u001B[39mn_rows))\n",
      "\u001B[0;31mAttributeError\u001B[0m: No variables to plot"
     ]
    }
   ],
   "source": [
    "dfn_models = [dfn_model, new_dfn_model]\n",
    "sims = []\n",
    "for model in dfn_models:\n",
    "    plot_sim = pybamm.Simulation(model)\n",
    "    plot_sim.solve([0, 3600])\n",
    "    sims.append(plot_sim)\n",
    "\n",
    "pybamm.dynamic_plot(sims, time_unit=\"seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To be able to plot the results from a serialised model, the mesh and model variables need to be saved alongside the model itself.\n",
    "\n",
    "To do this, set the `variables` option to `True` when saving the model as in the example below; notice how the models will now plot nicely."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "81d8329fab424264bd56c65d53d34f63",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(FloatSlider(value=0.0, description='t', max=3600.0, step=36.0), Output()), _dom_classes=…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<pybamm.plotting.quick_plot.QuickPlot at 0x111963010>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# using the first simulation, save a new file which includes a list of the model variables\n",
    "dfn_sim.save_model(\"sim_model_variables\", variables=True)\n",
    "\n",
    "# read the model back in\n",
    "model_with_vars = pybamm.load_model(\"sim_model_variables.json\")\n",
    "\n",
    "# plot the pre- and post-serialisation models together to prove they behave the same\n",
    "models = [dfn_model, model_with_vars]\n",
    "sims = []\n",
    "for model in models:\n",
    "    sim = pybamm.Simulation(model)\n",
    "    sim.solve([0, 3600])\n",
    "    sims.append(sim)\n",
    "\n",
    "pybamm.dynamic_plot(sims, time_unit=\"seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Saving from Model\n",
    "\n",
    "Alternatively, the model can be saved directly from the Model class.\n",
    "\n",
    "Note that at the moment, only models derived from the BaseBatteryModel class can be serialised; those built from scratch using pybamm.BaseModel() are currently unsupported.\n",
    "\n",
    "First set up the model, as explained in detail for the [SPM](https://github.com/pybamm-team/PyBaMM/blob/develop/docs/source/examples/notebooks/models/SPM.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<pybamm.models.full_battery_models.lithium_ion.spm.SPM at 0x29cf65c90>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# create the model\n",
    "spm_model = pybamm.lithium_ion.SPM()\n",
    "\n",
    "# set up and discretise ready to solve\n",
    "geometry = spm_model.default_geometry\n",
    "param = spm_model.default_parameter_values\n",
    "param.process_model(spm_model)\n",
    "param.process_geometry(geometry)\n",
    "mesh = pybamm.Mesh(geometry, spm_model.default_submesh_types, spm_model.default_var_pts)\n",
    "disc = pybamm.Discretisation(mesh, spm_model.default_spatial_methods)\n",
    "disc.process_model(spm_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then save the model. Note that in this case the model variables and the mesh must be provided directly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Serialise the spm_model, providing the varaibles and the mesh\n",
    "spm_model.save_model(\"example_model\", variables=spm_model.variables, mesh=mesh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now you can read the model back in, solve and plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ce5addf4f59c447e97d2fbee633cb6e0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(FloatSlider(value=0.0, description='t', max=1.0, step=0.01), Output()), _dom_classes=('w…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<pybamm.plotting.quick_plot.QuickPlot at 0x29c8a3d10>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# read back in\n",
    "new_spm_model = pybamm.load_model(\"example_model.json\")\n",
    "\n",
    "# select a solver and solve\n",
    "new_spm_solver = new_spm_model.default_solver\n",
    "new_spm_solution = new_spm_solver.solve(new_spm_model, [0, 3600])\n",
    "\n",
    "# plot the solution\n",
    "new_spm_solution.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Making edits to a serialised model\n",
    "\n",
    "As mentioned at the begining of this notebook, only models which have already been discretised can be serialised and readh back in. This means that after serialisation, the model *cannot be edited*, as the non-discretised elements of the model such as the original rhs are not saved.\n",
    "\n",
    "If you are likely to want to save a model and then edit it in the future, you may wish to use the `Simulation.save()` functionality to pickle your simulation, as described in [tutorial 6](https://github.com/pybamm-team/PyBaMM/blob/develop/docs/source/examples/notebooks/getting_started/tutorial-6-managing-simulation-outputs.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before finishing we will remove the data files we saved so that we leave the directory as we found it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.remove(\"example_model.json\")\n",
    "os.remove(\"sim_model_example.json\")\n",
    "os.remove(\"sim_model_variables.json\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "The relevant papers for this notebook are:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1] Joel A. E. Andersson, Joris Gillis, Greg Horn, James B. Rawlings, and Moritz Diehl. CasADi – A software framework for nonlinear optimization and optimal control. Mathematical Programming Computation, 11(1):1–36, 2019. doi:10.1007/s12532-018-0139-4.\n",
      "[2] Marc Doyle, Thomas F. Fuller, and John Newman. Modeling of galvanostatic charge and discharge of the lithium/polymer/insertion cell. Journal of the Electrochemical society, 140(6):1526–1533, 1993. doi:10.1149/1.2221597.\n",
      "[3] Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, and others. Array programming with NumPy. Nature, 585(7825):357–362, 2020. doi:10.1038/s41586-020-2649-2.\n",
      "[4] Scott G. Marquis, Valentin Sulzer, Robert Timms, Colin P. Please, and S. Jon Chapman. An asymptotic derivation of a single particle model with electrolyte. Journal of The Electrochemical Society, 166(15):A3693–A3706, 2019. doi:10.1149/2.0341915jes.\n",
      "[5] Valentin Sulzer, Scott G. Marquis, Robert Timms, Martin Robinson, and S. Jon Chapman. Python Battery Mathematical Modelling (PyBaMM). Journal of Open Research Software, 9(1):14, 2021. doi:10.5334/jors.309.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "pybamm.print_citations()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
